import numpy as np
import cmath 
import math 
import matplotlib.pyplot as plt
from scipy.fft import fft, fftfreq, rfft
import time

from dft_implement_1 import dft_transform

from tools import binary_to_decimal, decimal_to_binary, reverse_binary, reverse_order, get_rotation_factor

from pylab import mpl

mpl.rcParams['font.sans-serif'] = ['FangSong'] # 指定默认字体
mpl.rcParams['axes.unicode_minus'] = False # 解决保存图像是负号'-'显示为方块的问题


def fft_freq(N, d):
	fk_negative = np.arange(start=-N/2,stop=0,step=1)*(1/(N*d))
	fk_positive = np.arange(start=0,stop=N/2,step=1)*(1/(N*d))
	fk = np.concatenate((fk_positive, fk_negative))
	return fk 


def fft_transform(Xn, N=None):
	# if N == None:
	# 	N = len(Xn)
	# #print(save_array)
	# Xn_copy = reverse_order(Xn, N)
	if N == None:
		N = len(Xn)
	xx = Xn.copy()
	if math.modf(math.log2(N))[0] != 0:
		N_ = int(math.pow(2,math.modf(math.log2(N))[1]+1))
		print(N_, N)
		xx = np.pad(xx,(0,N_-N),mode="constant",constant_values=0)
		N = N_
	print("xx:",len(xx))
	Xn_copy = reverse_order(xx, N)

	rotation_factor = get_rotation_factor(N)
	sum = 0
	start = time.time()	# 开始计算
	L = int(math.log2(N))
	for m in range(L):
		m = m + 1	# 第m级蝶形运算
		k_sum = int(math.pow(2,m-1))
		#for count in range(N//int(math.pow(2,m))):	
		count = 0
		while count < N:
			# 第m级，蝶形运算总数
			#count = count + 1
			for k in range(k_sum):
				#t_1 = Xn_copy[2*count+k] + Xn_copy[2*count+k+k_sum]*rotation_factor[k<<(L-m)]
				#t_2 = Xn_copy[2*count+k] - Xn_copy[2*count+k+k_sum]*rotation_factor[k<<(L-m)]
				tt = Xn_copy[count+k+k_sum]*rotation_factor[k<<(L-m)]
				t_1 = Xn_copy[count+k] + tt
				start_2 = time.time()
				t_2 = Xn_copy[count+k] - tt
				Xn_copy[count+k+k_sum] = t_2
				sum += time.time() - start_2
				#print("print: ",t_1, t_2, count+k, count+k+k_sum, k<<(L-m), rotation_factor[k<<(L-m)])
				#Xn_copy[2*count+k] = t_1
				#Xn_copy[2*count+k+k_sum] = t_2
				#print("print: ",t_1, t_2, count+k, count+k+k_sum, Xn_copy[count+k], Xn_copy[count+k+k_sum])
				Xn_copy[count+k] = t_1
				Xn_copy[count+k+k_sum] = t_2
				
			#print(Xn_copy)
			count = count + int(math.pow(2,m))
	#print(Xn_copy)
	end = time.time()
	running_time = end - start
	print("sum:", sum)
	return Xn_copy, running_time

def fft_transform_real(Xn, N=None):
	if N == None:
		N = len(Xn)
	xx = Xn.copy()
	if math.modf(math.log2(N))[0] != 0:
		N_ = int(math.pow(2,math.modf(math.log2(N))[1]+1))
		print(N_, N)
		xx = np.pad(xx,(0,N_-N),mode="constant",constant_values=0)
		N = N_-N
	print("xx:",len(xx))
	xxx = np.ones(N, dtype=complex)
	Xn_copy = reverse_order(xx, N)
	rotation_factor = get_rotation_factor(N)
	pow_2_m = []
	for i in range(int(math.log2(N)+1)):
		pow_2_m.append(math.pow(2,i))

	sum = 0
	start = time.time()	# 开始计算
	L = int(math.log2(N))
	for m in range(L):
		m = m + 1	# 第m级蝶形运算
		k_sum = int(pow_2_m[m-1])
		#for count in range(N//int(math.pow(2,m))):	
		count = 0
		i = 0
		while count < N:
			# 第m级，蝶形运算总数
			#count = count + 1
			for k in range(k_sum):
				#print("copy:",Xn_copy)
				if k == 0:
					t_1 = Xn_copy[count+k] + Xn_copy[count+k+k_sum]*rotation_factor[k<<(L-m)]
					start_2 = time.time()
					t_2 = Xn_copy[count+k] - Xn_copy[count+k+k_sum]*rotation_factor[k<<(L-m)]
					xxx[count+k+k_sum] = t_2
					sum += time.time() - start_2
					#print("print: ",t_1, t_2, count+k, count+k+k_sum, Xn_copy[count+k], Xn_copy[count+k+k_sum])
					xxx[count+k] = t_1
					xxx[count+k+k_sum] = t_2
				else:
					# t_1 = Xn_copy[count+k] + Xn_copy[count+k+k_sum]*rotation_factor[k<<(L-m)]
					# t_2 = complex(t_1.real, -t_1.imag)
					# print("print: ",count+k, count+k+k_sum, Xn_copy[count+k], Xn_copy[count+k+k_sum])
					# Xn_copy[count+k] = t_1
					# Xn_copy[N-count-k] = t_2
					xxx[count+k] = Xn_copy[count+k] + Xn_copy[count+k+k_sum]*rotation_factor[k<<(L-m)]
					#xxx[int((math.pow(2,m)*(1+i)-k))] = complex(xxx[count+k].real, -xxx[count+k].imag)
					start_2 = time.time()
					#print(t_1)
					#xxx[count+k].conjugate();		
					xxx[int((pow_2_m[m]*(1+i)-k))] = complex(xxx[count+k].real, -xxx[count+k].imag)
					sum += time.time() - start_2

					#print("print: ",count+k, int((math.pow(2,m)-k)+math.pow(2,m)*i), count+k+k_sum, Xn_copy[count+k], Xn_copy[count+k+k_sum],
					#	xxx[count+k], xxx[int((math.pow(2,m)-k)+math.pow(2,m)*i)])
				#print("copy:",Xn_copy)
			i += 1		
			#print(Xn_copy)
			count = count + int(pow_2_m[m])
		Xn_copy = xxx.copy()
		#print("Xn_copy change")
	end = time.time()
	running_time = end - start
	#print(Xn_copy)
	print("real_sum:",sum)
	return Xn_copy, running_time


def generate_cos_wave(freq, sample_rate, duration):
    x = np.linspace(0, duration, sample_rate * duration, endpoint=False)
    frequencies = x * freq
    # 2pi because np.sin takes radians
    # 2*pi*f*n*T
    y = np.cos((2 * np.pi) * frequencies, dtype=complex)
    return x, y



def main():
	duration = 1
	freq_1 = 10
	freq_2 = 100
	freq_3 = 200
	p = 9

	sample_rate = int(math.pow(2,11))
	N = duration * sample_rate	
	x, xn_1 = generate_cos_wave(freq_1, sample_rate, duration)
	x, xn_2 = generate_cos_wave(freq_2, sample_rate, duration)
	x, xn_3 = generate_cos_wave(freq_3, sample_rate, duration)

	fk = fft_freq(sample_rate, 1/sample_rate)

	Xk_1, _ = fft_transform(xn_1, N)
	Xk_2, _ = fft_transform(xn_2, N)
	Xk_3, _ = fft_transform(xn_3, N)

	plt.figure(1)
	plt.subplot(321)
	plt.plot(x[:300],xn_1[:300])
	plt.title("xn_1(f=10HZ)")
	plt.subplot(322)
	plt.stem(fk,abs(Xk_1))
	plt.title("Xk_1")
	plt.subplot(323)
	plt.plot(x[:100],xn_2[:100])
	plt.title("xn_2(f=100HZ)")
	plt.subplot(324)
	plt.stem(fk,abs(Xk_2))
	plt.title("Xk_2")
	plt.subplot(325)
	plt.plot(x[:60],xn_3[:60])
	plt.title("xn_3(f=200HZ)")
	plt.subplot(326)
	plt.stem(fk,abs(Xk_3))
	plt.title("Xk_3")
	plt.show()

	xn = xn_1+xn_2+xn_3
	Xk, running_time = fft_transform(xn,N)
	#Xk_real, real_running_time = fft_transform_real(xn, N)
	plt.subplot(2,1,1)
	plt.plot(x[:200],xn[:200])
	plt.title("Xn(混合信号)")
	plt.subplot(2,1,2)
	
	plt.stem(fk,abs(Xk))
	plt.title("Xk")
	plt.show()

def main2():
	duration = 1
	freq_1 = 10
	freq_2 = 100
	freq_3 = 200
	p = 9

	#sample_rate = int(math.pow(2,p))
	fs = 128

	N = duration * fs
	x, xn_1 = generate_cos_wave(freq_1, fs, duration)
	x, xn_2 = generate_cos_wave(freq_2, fs, duration)
	x, xn_3 = generate_cos_wave(freq_3, fs, duration)
	xn = xn_1+xn_2+xn_3

	fk = fft_freq(N, 1/fs)

	Xk, _ = fft_transform(xn,N)

	plt.figure(1);
	plt.subplot(321);
	plt.title("fs=128HZ")
	plt.xlabel("f/HZ")
	plt.ylabel("|Xk|")
	print(len(fk),len(Xk))
	plt.stem(fk, abs(Xk))
	plt.subplot(322)
	ph = np.angle(Xk)
	plt.plot(fk, ph)



	fs = 256

	N = duration * fs
	x, xn_1 = generate_cos_wave(freq_1, fs, duration)
	x, xn_2 = generate_cos_wave(freq_2, fs, duration)
	x, xn_3 = generate_cos_wave(freq_3, fs, duration)
	xn = xn_1+xn_2+xn_3

	fk = fft_freq(256, 1/fs)

	Xk, _ = fft_transform(xn,N)

	plt.subplot(323);
	plt.title("fs=256HZ")
	plt.xlabel("f/HZ")
	plt.ylabel("|Xk|")
	plt.stem(fk, abs(Xk))
	plt.subplot(324)
	ph = np.angle(Xk)
	plt.plot(fk, ph)

	fs = 512

	N = duration * fs
	x, xn_1 = generate_cos_wave(freq_1, fs, duration)
	x, xn_2 = generate_cos_wave(freq_2, fs, duration)
	x, xn_3 = generate_cos_wave(freq_3, fs, duration)
	xn = xn_1+xn_2+xn_3

	fk = fft_freq(N, 1/fs)

	Xk, _ = fft_transform(xn,N)

	print(len(fk))
	print(Xk)

	plt.subplot(325);
	plt.title("fs=512HZ")
	plt.xlabel("f/HZ")
	plt.ylabel("|Xk|")
	plt.stem(fk, abs(Xk))
	print(max(fk))
	plt.subplot(326)
	ph = np.angle(Xk)
	plt.plot(fk, ph)

	plt.show()


	fs = 600

	N = duration * fs
	x, xn_1 = generate_cos_wave(freq_1, fs, duration)
	x, xn_2 = generate_cos_wave(freq_2, fs, duration)
	x, xn_3 = generate_cos_wave(freq_3, fs, duration)
	xn = xn_1+xn_2+xn_3

	fk = fft_freq(1024, 1/fs)

	Xk, _ = fft_transform(xn,N)

	print(len(fk))
	print(len(Xk))

	plt.stem(fk, abs(Xk))

	plt.show()




if __name__ == '__main__':
	main2()